Description

Function reshape() from modules mxnet.numpy and mxnet.numpy_extension will now be executed by oneDNN primitives as it is done in module mxnet.ndarray.

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage

Changes

• Add oneDNN execution attributes to NNVM_REGISTER_OP(_npx_reshape)

Comments

Using two NNVM_REGISTER_OP() functions (one for mxnet.ndarray, second for mxnet.numpy/mxnet.numpy_extension) instead of one with .add_alias() because special values in npx.reshape's and nd.reshape's new_shape/target_shape parameters have different functionality. For example:

in nd.reshape():

• "-3" use the product of two consecutive dimensions of the input shape as the output dimension.

in npx.reshape():

• "-3" will skip current dimension if and only if the current dim size is one.

Performance comparison:

• ~50% peformance speedup for mx.np.reshape() and mx.npx.reshape().
• Slight loss of performance for mx.nd.reshape() is caused by adding 0-shape checks to oneDNN path (necessary for mx.np and mx.npx). It is a big deal only for small tensors.